;;; shock tube problem solver
;;; (c) http://www.grc.nasa.gov/WWW/wind/valid/stube/stube.html

(import (rnrs)
        (grid-gen common)
)

(module-on fortran77)

(define grid-dims '(1000))

(define phys-q (comp-vector (comp-float "rho")
                            (comp-float "p")
                            (comp-float "v")
                            (comp-float "x")
))

(define grid (comp-array phys-q grid-dims))

(define xmin (comp-float "xmin"  0.0))
(define xmax (comp-float "xmax"  1.0))
(define xmid (comp-float "x0"  0.5))

(define dt (comp-float "dt" 0.18))

(define rho4 (comp-float "rho4"  1.0))
(define p4   (comp-float "p4"    3.0))
(define v4   (comp-float "v4"    0.0))

(define rho1 (comp-float "rho1"  1.0))
(define p1   (comp-float "p1"    1.0))
(define v1   (comp-float "v1"    0.0))

(define rho2 (comp-float "rho2"))
(define p2 (comp-float "p2"))
(define v2 (comp-float "v2"))

(define rho3 (comp-float "rho3"))
(define p3 (comp-float "p3"))
(define v3 (comp-float "v3"))

(define gam (comp-float "gam" '(/ 5. 3.)))

(define gm1 `(- ,gam 1.))
(define gp1 `(+ ,gam 1.))

(define cson (lambda (p rho) `(sqrt (/ (* ,gam ,p) ,rho))))

(define c1 (comp-float "c1"))
(define c2 (comp-float "c2"))
(define c3 (comp-float "c3"))
(define c4 (comp-float "c4"))

(define cod (code-gen "shocktube"
  (beg
    (set-val c1 (cson p1 rho1))
    (set-val c4 (cson p4 rho4))

    (define dl (comp-float "dl"))
    (define p2p1m (comp-float "p2p1m"))
    (define p2p1n (comp-float "p2p1n"))
    (define p2p1 (comp-float "p2p1"))
    (define t1 (comp-float "t1"))
    (define t2 (comp-float "t2"))
    (define t3 (comp-float "t3"))
    (define fm (comp-float "fm"))
    (define f  (comp-float "f"))
    
    (set-val p2p1m (* 0.9 (/ p4 p1)))
    (set-val t1 (/ (* -2.0 gam) gm1))
    (set-val t2 (* gm1
                   (/ c1 c4)
                   (- p2p1m 1.0)
                )
    )
    (set-val t3 (* 2.0
                   gam
                   (+ (* 2.0 gam)
                      (* gp1 (- p2p1m 1.0))
                   )
                )
    )
    (set-val fm (- (/ p4 p1)
                   (* p2p1m
                      (** (- 1.0
                             (/ t2 (sqrt t3))
                          )
                          t1
                      )
                   )
                )
    )
    (set-val p2p1 (* 0.95 p2p1m))
    
    (set-val f 1.0)

    (do-while (> (abs f) 1.0e-10)
      (beg
        (set-val t2 (* gm1 (/ c1 c4) (- p2p1 1.0)))
        (set-val t3 (* 2.0 gam (+ (* 2.0 gam) (* gp1 (- p2p1 1.0)))))
        (set-val f  (- (/ p4 p1) (* p2p1 (** (- 1.0 (/ t2 (sqrt t3))) t1))))
        
        (print "p2p1,f:" p2p1 f)
        
        (set-val p2p1n (- p2p1 (* f (/ (- p2p1 p2p1m) (- f fm)))))
        (set-val p2p1m p2p1)
        (set-val fm f)
        (set-val p2p1 p2p1n)
      )
    )
    
    (define t2t1 (comp-float "t2t1"))
    (define r2r1 (comp-float "r2r1"))
    (define wsp (comp-float "wsp"))
    (define xs (comp-float "xs"))
    (define xc (comp-float "xc"))
    (define xh (comp-float "xh"))
    (define xt (comp-float "xt"))
    (define up (comp-float "up"))
    
    (set-val t2t1 (* p2p1 
                     (/ (+ (/ gp1 gm1) p2p1)
                        (+ 1.0 (* gp1 (/ p2p1 gm1)))
                     )
                  )
    )
    (set-val r2r1 (/ (+ 1.0 (* gp1 (/ p2p1 gm1)))
                     (+ (/ gp1 gm1) p2p1)
                  )
    )
    (set-val wsp (* c1
                    (sqrt (+ (* gp1
                                (/ (- p2p1 1.0)
                                   (* 2.0 gam)
                                )
                             )
                             1.0
                          )
                    )
                 )
    )
    (set-val xs (+ xmid (* wsp dt)))
    (set-val p2 (* p2p1 p1))
    (set-val rho2 (* r2r1 rho1))
    (set-val p3 p2)
    (set-val rho3 (* rho4 (** (/ p3 p4) (/ 1.0 gam))))
    (set-val c3 (cson p3 rho3))
    (set-val up (/ (* 2.0
                      c4
                      (- 1.0
                         (** (/ p2 p4)
                             (* 0.5 (/ gm1 gam))
                         )
                      )
                   )
                   gm1
                )
    )
    (set-val xc (+ xmid (* up dt)))
    (set-val xh (+ xmid (* (- v4 c4) dt)))
    (set-val xt (+ xmid (* (- up c3) dt)))

    (arr-loop grid
      (let ((x   (set-ind (vec-var grid "x")   (cur-ptr)))
            (rho (set-ind (vec-var grid "rho") (cur-ptr)))
            (p   (set-ind (vec-var grid "p")   (cur-ptr)))
            (v   (set-ind (vec-var grid "v")   (cur-ptr)))
            (i   (cur-ptr 0))
           )
        (set-val x (+ xmin 
                      (* (- xmax xmin) 
                         (/ (- i 1.0)
                            (- (car grid-dims) 1.0)
                         )
                      )
                   )
        )
        (if (<= x xh)
            (beg
              (set-val rho rho4)
              (set-val p p4)
              (set-val v v4)
            )
            (if (<= x xt)
                (beg
                  (set-val v (+ v4
                                (* up
                                   (/ (- x xh) (- xt xh))
                                )
                             )
                  )
                  (set-val p (* p4
                                (** (- 1.0
                                       (* 0.5 gm1 (/ v c4))
                                    )
                                    (* 2.0 (/ gam gm1))
                                )
                             )
                  )
                  (set-val rho (* rho4
                                  (** (- 1.0
                                         (* 0.5 gm1 (/ v c4))
                                      )
                                      (/ 2.0 gm1)
                                  )
                               )
                  )
                )
                (if (< x xc)
                    (beg
                      (set-val rho rho3)
                      (set-val p p3)
                      (set-val v up)
                    )
                    (if (< x xs)
                        (beg
                          (set-val rho rho2)
                          (set-val p p2)
                          (set-val v up)
                        )
                        (beg
                          (set-val rho rho1)
                          (set-val p p1)
                          (set-val v v1)
                        )
                    )
                )
            )
        )
      )
    )
    (let ((x   (vec-var grid "x"))
          (rho (vec-var grid "rho"))
          (p   (vec-var grid "p"))
          (v   (vec-var grid "v"))
         )
      (save-txt "rho.dat" x rho)
      (save-txt "p.dat" x p)
      (save-txt "v.dat" x v)
    )
  )
))
